from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
import pickle
from pathlib import Path
from langchain.schema import Document
from langchain.retrievers import BM25Retriever
from reranking import LLMReranker

class VectorRetriever:
    def __init__(self, vector_db_path: str, checkpoint: str = "./jinaai/jina-embeddings-v2-base-zh"):
        self.vector_db_path = vector_db_path
        
        model_name = checkpoint
        model_kwargs = {"device": "cpu"}
        encode_kwargs = {"normalize_embeddings": False}

        embeddings = HuggingFaceEmbeddings(
            model_name=model_name,
            model_kwargs=model_kwargs,
            encode_kwargs=encode_kwargs,
        )
        self.embeddings = embeddings

        # Load the FAISS vector store
        faiss_retriever = FAISS.load_local(
            folder_path=self.vector_db_path,
            embeddings=self.embeddings,
            allow_dangerous_deserialization=True
        )
        self.faiss_retriever = faiss_retriever

    def get_relevant_documents(self, query: str):
        # Perform similarity search using FAISS
        docs = self.faiss_retriever.as_retriever().get_relevant_documents(query)

        # filtered_docs = [doc for doc in docs if doc.metadata.get("institution") == "人行"]
        return docs
    def get_relevant_documents_with_score(self, query: str, top_k: int = 30):
        if self.faiss_retriever is None:
            raise ValueError("FAISS vector store not loaded. Call load_vector_db() first.")
        # Perform similarity search with scores using FAISS
        docs, scores = zip(*self.faiss_retriever.similarity_search_with_score(query, k=top_k))
        for doc, score in zip(docs, scores):
            doc.metadata["fassi_score"] = score

        return docs
        

class bm25Retriever:
    def __init__(self, bm25_index_path: str):
        self.bm25_index_path = bm25_index_path
        with open(bm25_index_path, 'rb') as f:
            bm25_retriever: BM25Retriever = pickle.load(f)
        self.bm25_retriever = bm25_retriever


    def get_relevant_documents(self, query: str, top_k: int = 4):
        if self.bm25_retriever is None:
            raise ValueError("BM25 index not loaded. Call load_bm25_index() first.")
        
        return self.bm25_retriever.get_relevant_documents(query, top_k=top_k)
        

class HybridRetriever:
    def __init__(self, vector_db_path: str, bm25_index_path: str):
        self.vector_db = VectorRetriever(vector_db_path)
        self.reranker = LLMReranker()
    
    def get_relevant_documents_with_score(self, query: str,
                                          llm_reranking_sample_size: int = 20,
                                        documents_batch_size=3,
                                        llm_weight: float = 0.7,
                                        top_k: int = 4):
        # Get documents from the vector database
        vector_docs = self.vector_db.get_relevant_documents_with_score(query, top_k=llm_reranking_sample_size)
        reranked_results = self.reranker.rerank_documents(
            query=query,
            documents=vector_docs,
            documents_batch_size=documents_batch_size,
            llm_weight=llm_weight,
            top_k=top_k
        )

